import argparse
import math
import sys
import numpy as np
import cv2
import multiprocessing

sys.path.append("/home/nuc2/PycharmProjects/yolov5-master")
from utils_track import vino as ov
from Map_Reflect_utils import Realsense
import time

redcode = 0
bluecode = 1


class HSVTrackball:
    def __init__(self, color_mode=redcode):
        self.min_coord = (750, 300)
        self.share_data = dict  # 共享字典
        self.color_mode = color_mode  # 控制球的颜色

    def color_track_main(self, model):
        realsense_cam = Realsense.realsense(1)
        cap = realsense_cam.cam_init(1080)
        while True:
            start = time.time()
            min_distance = np.inf
            frame, depth_colormap, depth_intrin, aligned_depth_frame = realsense_cam.cam_run(cap)
            frame, det = model.run(frame)
            self.min_coord = (750, 340)
            if len(det):
                for xyxy, conf, cls in reversed(det):  # 识别物的类别像素坐标，置信度，物体类别号
                    mid_pos = [int((xyxy[0] + xyxy[2]) / 2),
                               int((xyxy[1] + xyxy[3]) / 2)]
                    min_val = min(abs(int(xyxy[2]) - int(xyxy[0])), abs(int(xyxy[3]) - int(xyxy[1])))
                    # 将中心点向四周扩展计算深度

                    z, y, x = realsense_cam.depth_to_data(depth_intrin, aligned_depth_frame, (750,340),
                                                          min_val)
                    if min_distance > math.sqrt(
                            abs(self.min_coord[0] - mid_pos[0]) ** 2 + abs(self.min_coord[1] - mid_pos[1]) ** 2):
                        min_distance = math.sqrt(
                            abs(self.min_coord[0] - mid_pos[0]) ** 2 + abs(self.min_coord[1] - mid_pos[1]) ** 2)

                        self.min_coord = (int(mid_pos[0]), int(mid_pos[1]), round(x, 4))
                        print(x)
                    cv2.circle(frame, mid_pos, 1, (255, 0, 0), 3)  # 圆心
                    cv2.putText(frame, f"{mid_pos}", mid_pos, cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 2)
                    # 由于识别的时候将画面转移了90°，为了对齐深度的像素，需要将像素逆变换回去，才能拿

            end = time.time() if time.time() != start else start + 0.01
            fps = 1 / (end - start)
            cv2.putText(frame, f"{round(fps, 1)}", (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1,
                        (255, 0, 0),
                        2)
            # print(self.share_data["coord_data"])
            # self.share_data["im0"] = color_image_Processed_upper
            self.share_data["coord_data"] = self.min_coord
            self.share_data["im0"] = frame

    # This function is used to choose argument,which is called init function
    def parse_opt(self):
        parser = argparse.ArgumentParser()
        parser.add_argument('--model_path', type=str,
                            default="/home/nuc2/PycharmProjects/yolov5-master/weights/best_ball.xml",
                            help='model path')
        parser.add_argument('--weights_path', nargs='+', type=str,
                            default="/home/nuc2/PycharmProjects/yolov5-master/weights/best_ball.bin",
                            help='weights path or triton URL')
        parser.add_argument('--conf-thres', type=float, default=0.01, help='confidence threshold')
        parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
        parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
        parser.add_argument('--classes', type=list, default=[self.color_mode], help='Classes')
        parser.add_argument('--img-size', type=int, default=640, help='img-size')
        parser.add_argument('--device', default='GPU', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
        parser.add_argument('--name-porcess', default='BucketTime', help='name to processs--cjh')
        parser.add_argument('--source', type=str, default="0",
                            help='')
        opt = parser.parse_args()

        return opt

    def track_init(self, shareVar):
        self.share_data = shareVar
        # 使其为全局变量，不然无法添加
        opt = self.parse_opt()
        del vars(opt)["source"]
        model = ov.Vino(**vars(opt))
        self.color_track_main(model)


if __name__ == '__main__':
    shareVar = multiprocessing.Manager().dict()
    track = HSVTrackball(redcode)
    p3 = multiprocessing.Process(target=track.track_init, args=(shareVar,))
    p3.start()
    while True:
        try:
            print(shareVar["coord_data"])
        # cv2.imshow("1231", shareVar['im0'])
        # cv2.imshow("1232", shareVar['im1'])
        # cv2.waitKey(1)
        except Exception:
            pass
